from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from database import get_db
from models import User
from passlib.context import CryptContext
from jose import jwt
from datetime import datetime, timedelta
from pydantic import BaseModel
from typing import Union
from jose import JWTError
from fastapi.security import OAuth2PasswordBearer
# 初始化 oauth2_scheme
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")


router = APIRouter()

# 用来加密/验证密码
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# 生成 token 的配置
SECRET_KEY = "3a1f5f1d6c6a9d7c8b2a9f6d91c0e4f54e0f43f9d1b823c3a36db5a4e9f63e3a"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60

# 登录请求体
class LoginRequest(BaseModel):
    username: str
    password: str

# 登录返回体
class LoginResponse(BaseModel):
    access_token: str
    token_type: str
    role: str


def create_access_token(data: dict, expires_delta: Union[timedelta, None] = None):
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

# 关键：获取当前用户（之前丢失的函数）
def get_current_user(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
    credentials_exception = HTTPException(
        status_code=401,
        detail="认证失败",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
    except JWTError:
        raise credentials_exception
    user = db.query(User).filter(User.username == username).first()
    if user is None:
        raise credentials_exception
    return user


@router.post("/login", response_model=LoginResponse)
def login(req: LoginRequest, db: Session = Depends(get_db)):
    user = db.query(User).filter(User.username == req.username).first()
    if not user:
        raise HTTPException(status_code=401, detail="用户不存在")

    if not pwd_context.verify(req.password, user.password):
        raise HTTPException(status_code=401, detail="密码错误")

    # 生成 JWT token
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username, "role": user.role},
        expires_delta=access_token_expires
    )

    return {
        "access_token": access_token,
        "token_type": "bearer",
        "role": user.role
    }
